
#include "histogram_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <cstdio>
const uint32_t IO_CHANNELS = 16;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  HistogramTilingData tiling;
  const gert::StorageShape* x1_shape = context->GetInputShape(0);
  int32_t totalLen = context->GetInputShape(0)->GetOriginShape().GetShapeSize();
  const gert::RuntimeAttrs *attrs = context->GetAttrs();
  int32_t bins = *(attrs->GetAttrPointer<int32_t>(0));
  float min = *(attrs->GetAttrPointer<float>(1));
  float max = *(attrs->GetAttrPointer<float>(2));
  
  int32_t dtypeSizeX = ge::GetSizeByDataType(context->GetInputDesc(0)->GetDataType());
  int32_t blockLen = 32/dtypeSizeX;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  uint64_t ubSize = 0; // ub剩余空间大小
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
  int32_t bufferLen = ubSize/IO_CHANNELS/sizeof(float);
  bufferLen -= bufferLen%blockLen;
  int32_t groups = totalLen/bufferLen;
  int32_t tailLen = totalLen-bufferLen*groups;
  int32_t tailLenPadded = (tailLen+blockLen-1)/blockLen*blockLen;
  
  printf("totalLen: %d, bins: %d, min: %f, max: %f\n", totalLen, bins, min, max);
  if(min==0 && max==0){
    printf("calc min max!\n");
    const gert::Tensor *x_tensor = context->GetInputTensor(0);
    switch(x_tensor->GetDataType()){
        case ge::DT_FLOAT:{
            auto dt = x_tensor->GetData<float>();
            min = 1e38;
            max = -1e38;
            for(int32_t i=0;i<totalLen;i++){
                min = dt[i]<min?dt[i]:min;
                max = dt[i]>max?dt[i]:max;
            }
            break;
        }
        case ge::DT_FLOAT16:{
            auto dt = x_tensor->GetData<int16_t>();
            min = 1e38;
            max = -1e38;
            for(int32_t i=0;i<totalLen;i++){
                int16_t base = dt[i]&0x03ff;
                int16_t pow = ((dt[i]>>10)&0x001f)-15;
                float res = base/1024.0+1;
                if(pow<0)res/=((static_cast<uint32_t>(1))<<(-pow));
                else if(pow>0)res*=((static_cast<uint32_t>(1))<<(pow));
                if(dt[i]<0)res = -res;
                min = res<min?res:min;
                max = res>max?res:max;
            }
            break;
        }
        case ge::DT_INT32:{
            auto dt = x_tensor->GetData<int32_t>();
            int32_t intMin = 0x7fffffff;
            int32_t intMax = 0x80000000;
            for(int32_t i=0;i<totalLen;i++){
                intMin = dt[i]<intMin?dt[i]:intMin;
                intMax = dt[i]>intMax?dt[i]:intMax;
            }
            min = static_cast<float>(intMin);
            max = static_cast<float>(intMax);
            break;
        }
    }
  }
  float offset = -min;
  float width = (max-min)/static_cast<float>(bins);

  tiling.set_totalLen(totalLen);
  tiling.set_bins(bins);
  tiling.set_min(min);
  tiling.set_max(max);
  tiling.set_offset(offset);
  tiling.set_width(width);
  tiling.set_blockLen(blockLen);
  tiling.set_bufferLen(bufferLen);
  tiling.set_groups(groups);
  tiling.set_tailLen(tailLen);
  tiling.set_tailLenPadded(tailLenPadded);

  printf("totalLen: %d, bins: %d, min: %f, max: %f\n", totalLen, bins, min, max);
  
  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    const gert::RuntimeAttrs *attrs = context->GetAttrs();
    int32_t bins = *(attrs->GetAttrPointer<int32_t>(0));
    y_shape->SetDimNum(1);
    y_shape->SetDim(0, bins);
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Histogram : public OpDef {
public:
    explicit Histogram(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("bins").AttrType(OPTIONAL).Int(100);
        this->Attr("min").AttrType(OPTIONAL).Float(0);
        this->Attr("max").AttrType(OPTIONAL).Float(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Histogram);
}
